"""
根据神经网络层搭建神经网络模型
"""

import dezero
import dezero.layers as L
import dezero.functions as F

class Model(L.Layer):
    def plot(self, *inputs, to_file='model.png'):
        y = self.forward(*inputs)
        return dezero.utils.plot_dot_graph(y,verbose=True,to_file=to_file)

class MLP(Model):
    def __init__(self, fc_output_sizes, activation=F.sigmoid):
        """
        名称：通用多层神经网络
        参数：
            fc_output_sizes：元组或列表，指定全连接层的每层输出大小
            activation：激活函数，默认sigmoid
        """
        super().__init__()
        self.activation = activation
        self.layers = []

        for i,out_size in enumerate(fc_output_sizes):
            layer = L.Linear(out_size)
            setattr(self, 'l'+str(i), layer) #给自己添加属性键值对如l1=layer1
            self.layers.append(layer)

    def forward(self,x):
        for l in self.layers[:-1]:
            x = self.activation(l(x))
        return self.layers[-1](x) #最后一层是输出层，输出层的激活函数和隐藏层的激活函数不同，此处回归分析一般使用恒等函数即直接输出


#==========================
# 代表性的CNN  VGG16
#==========================

class VGG16(Model):
    """
    卷积层1（2*conv3-64）,卷积层2（2*conv3-128）,卷积层3（3*conv3-256）,卷积层4（3*conv3-512）,卷积层5（3*conv3-512）
    在每两层之间有池化层为移动步长为2的2X2池化矩阵（maxpool）。
    在卷积层5后有3个全连接层，再之后是soft-max预测层
    """

    WEIGHTS_PATH = 'https://github.com/koki0702/dezero-models/releases/download/v0.1/vgg16.npz'

    def __init__(self, pretrained=False):
        super().__init__()
        #只指定输出的通道数
        self.conv1_1 = L.Conv2d(64, kernel_size=3, stride=1, pad=1 )
        self.conv1_2 = L.Conv2d(64, kernel_size=3, stride=1, pad=1 )
        self.conv2_1 = L.Conv2d(128, kernel_size=3, stride=1, pad=1 )
        self.conv2_2 = L.Conv2d(128, kernel_size=3, stride=1, pad=1 )
        self.conv3_1 = L.Conv2d(256, kernel_size=3, stride=1, pad=1 )
        self.conv3_2 = L.Conv2d(256, kernel_size=3, stride=1, pad=1 )
        self.conv3_3 = L.Conv2d(256, kernel_size=3, stride=1, pad=1 )
        self.conv4_1 = L.Conv2d(512, kernel_size=3, stride=1, pad=1 )
        self.conv4_2 = L.Conv2d(512, kernel_size=3, stride=1, pad=1 )
        self.conv4_3 = L.Conv2d(512, kernel_size=3, stride=1, pad=1 )
        self.conv5_1 = L.Conv2d(512, kernel_size=3, stride=1, pad=1 )
        self.conv5_2 = L.Conv2d(512, kernel_size=3, stride=1, pad=1 )
        self.conv5_3 = L.Conv2d(512, kernel_size=3, stride=1, pad=1 )
        #只指定输出大小
        self.fc6 = L.Linear(4096)
        self.fc7 = L.Linear(4096)
        self.fc8 = L.Linear(1000)

        """如果需要加载已训练权重，则使用get_file获取网上的权重到本地使用"""
        if pretrained:
            weights_path = dezero.utils.get_file(VGG16.WEIGHTS_PATH)
            self.load_weights(weights_path)

    def forward(self,x):
        x = F.relu(self.conv1_1(x))
        x = F.relu(self.conv1_2(x))
        x = F.pooling(x,2,2)
        x = F.relu(self.conv2_1(x))
        x = F.relu(self.conv2_2(x))
        x = F.pooling(x,2,2)
        x = F.relu(self.conv3_1(x))
        x = F.relu(self.conv3_2(x))
        x = F.relu(self.conv3_3(x))
        x = F.pooling(x,2,2)
        x = F.relu(self.conv4_1(x))
        x = F.relu(self.conv4_2(x))
        x = F.relu(self.conv4_3(x))
        x = F.pooling(x,2,2)
        x = F.relu(self.conv5_1(x))
        x = F.relu(self.conv5_2(x))
        x = F.relu(self.conv5_3(x))
        x = F.pooling(x,2,2)

        x = F.reshape(x, (x.shape[0],-1)) #卷积层切换到全连接层的变形

        x = F.dropout( F.relu( self.fc6(x) ) )
        x = F.dropout( F.relu( self.fc7(x) ) )
        x = self.fc8(x)

        return x



